import os
import pandas as pd
import jieba.analyse
import json
from wordcloud import WordCloud
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt

# 中文显示
plt.rcParams['font.sans-serif'] = ['simsun']
plt.rcParams['axes.unicode_minus'] = False

# 停用词列表
stop_words = set([
    "的", "了", "是", "在", "与", "有", "为", "就", "和", "不", "很", "也",
    "这", "它", "我们", "你", "我", "出", "上", "还是", "这个", "之前", "app", "一直",
    "一款", "而且", "没有", "现在"
])

# 文件路径设置
input_folder = 'D:/Code/PycharmProjects/250611_clean/cleaned_data'
output_json = 'keywords_summary.json'
output_wordcloud_folder = 'D:/Code/PycharmProjects/250611_clean/wordclouds'
mask_img_path = 'D:/Code/PycharmProjects/250611_clean/photo/jdicon.jpg'
font_path = 'C:/Windows/Fonts/simsun.ttc'

# 创建词云输出目录（如果不存在）
os.makedirs(output_wordcloud_folder, exist_ok=True)

# 读取词云形状图
mask = np.array(Image.open(mask_img_path))

# 总关键词字典
per_file_keywords = {}

# 遍历所有CSV文件
for filename in os.listdir(input_folder):
    if filename.endswith('.csv'):
        file_path = os.path.join(input_folder, filename)
        df = pd.read_csv(file_path)

        # 提取关键词
        keywords = []
        for comment in df['评论']:
            if isinstance(comment, str):
                words = jieba.analyse.textrank(comment, topK=10, withWeight=False)
                words = [word for word in words if word not in stop_words]
                keywords.extend(words)

        # 保存到总关键词字典中
        per_file_keywords[filename] = list(set(keywords))

        # 生成词云
        wordcloud = WordCloud(
            font_path=font_path,
            background_color='white',
            max_words=300,
            mask=mask,
            scale=4,
            max_font_size=50,
            random_state=42,
            collocations=False
        ).generate(" ".join(keywords))

        # 显示词云
        plt.figure(figsize=(10, 6))
        plt.imshow(wordcloud, interpolation='bilinear')
        plt.axis("off")
        plt.tight_layout()
        plt.close()

        # 保存词云图像
        wc_filename = os.path.splitext(filename)[0] + '_wordcloud.png'
        wc_path = os.path.join(output_wordcloud_folder, wc_filename)
        wordcloud.to_file(wc_path)

# 保存关键词 JSON
with open(output_json, 'w', encoding='utf-8') as f:
    json.dump(per_file_keywords, f, ensure_ascii=False, indent=4)

print("所有词云生成完成，结果已保存在 wordclouds 文件夹；关键词保存为 keywords_summary.json。")
